import re
from dataclasses import dataclass
from typing import List, Optional, Tuple

@dataclass
class RustDeclaration:
    kind: str
    name: str
    visibility: str
    generics: Optional[str]
    line_number: int
    content: str

class RustParser:
    def __init__(self, source_code: str):
        self.source_code = source_code
        self.lines = source_code.split('\n')

    def parse_all(self) -> List[RustDeclaration]:
        declarations = []
        declarations.extend(self.find_functions())
        declarations.extend(self.find_structs())
        declarations.extend(self.find_unions())
        declarations.extend(self.find_enums())
        declarations.extend(self.find_statics())
        return declarations

    def get_complete_declaration(self, start_line_idx: int) -> Tuple[str, int]:
        """获取完整的声明，包括多行内容"""
        buffer = []
        brace_count = 0
        line_count = len(self.lines)
        current_idx = start_line_idx

        while current_idx < line_count:
            line = self.lines[current_idx].strip()

            if not line or line.startswith('//'):
                current_idx += 1
                continue

            buffer.append(line)

            # 计算花括号配对
            brace_count += line.count('{') - line.count('}')

            # 如果找到分号且不在宏内，或者括号配对完成
            if (line.endswith(';') and brace_count == 0) or (brace_count == 0 and brace_count != 0):
                break

            current_idx += 1

        return ' '.join(buffer), current_idx

    def find_macro_content(self, start_line_idx: int) -> Tuple[str, int]:
        """专门用于获取宏内容的方法"""
        buffer = []
        brace_count = 0
        line_count = len(self.lines)
        current_idx = start_line_idx
        started = False

        while current_idx < line_count:
            line = self.lines[current_idx].strip()

            if not line or line.startswith('//'):
                current_idx += 1
                continue

            # 检查宏的开始
            if '{' in line:
                started = True

            if started:
                buffer.append(line)
                brace_count += line.count('{') - line.count('}')

                if brace_count == 0:
                    break
            else:
                buffer.append(line)

            current_idx += 1

        return '\n'.join(buffer), current_idx

    def find_functions(self) -> List[RustDeclaration]:
        """查找所有函数声明，支持多行"""
        declarations = []
        i = 0
        while i < len(self.lines):
            line = self.lines[i].strip()

            # 跳过空行和注释
            if not line or line.startswith('//'):
                i += 1
                continue

            # 检查是否是函数声明的开始
            fn_pattern = r'(?P<vis>pub\s+)?(?:async\s+)?(?:unsafe\s+)?fn\s+(?P<name>[a-zA-Z_][a-zA-Z0-9_]*)'
            if re.match(fn_pattern, line) or (i > 0 and re.match(fn_pattern, self.lines[i-1].strip())):
                # 获取完整的声明
                full_decl, end_line = self.get_complete_declaration(i)

                # 提取函数名和可见性
                match = re.search(fn_pattern, full_decl)
                if match:
                    # 找到函数声明结束的位置（大括号之前）
                    decl_end = full_decl.find('{')
                    if decl_end == -1:
                        decl_end = len(full_decl)

                    # 提取声明部分
                    content = full_decl[:decl_end].strip()
                    if not content.endswith(';'):
                        content += ';'

                    # 构建声明对象
                    declarations.append(RustDeclaration(
                        kind="function",
                        name=match.group('name'),
                        visibility=match.group('vis').strip() if match.group('vis') else '',
                        generics=None,  # TODO: 添加泛型支持
                        line_number=i + 1,
                        content=content
                    ))

                i = end_line + 1
                continue

            i += 1

        return declarations

    def find_structs(self) -> List[RustDeclaration]:
        """查找所有结构体声明"""
        pattern = r'(?P<vis>pub\s+)?struct\s+(?P<name>[a-zA-Z_][a-zA-Z0-9_]*)\s*(?P<generics><[^>]*>)?'
        return self._find_declarations(pattern, "struct")

    def find_unions(self) -> List[RustDeclaration]:
        """查找所有联合体声明"""
        pattern = r'(?P<vis>pub\s+)?union\s+(?P<name>[a-zA-Z_][a-zA-Z0-9_]*)\s*(?P<generics><[^>]*>)?'
        return self._find_declarations(pattern, "union")

    def find_enums(self) -> List[RustDeclaration]:
        """查找所有枚举声明"""
        pattern = r'(?P<vis>pub\s+)?enum\s+(?P<name>[a-zA-Z_][a-zA-Z0-9_]*)\s*(?P<generics><[^>]*>)?'
        return self._find_declarations(pattern, "enum")

    def find_statics(self) -> List[RustDeclaration]:
        """查找所有静态变量声明，包括 lazy_static 宏中的声明"""
        declarations = []
        i = 0

        while i < len(self.lines):
            line = self.lines[i].strip()

            if not line or line.startswith('//'):
                i += 1
                continue

            # 处理 lazy_static 宏
            if line.startswith('lazy_static!'):
                macro_content, end_line = self.find_macro_content(i)
                # 在宏内容中查找静态声明
                static_pattern = r'(?P<vis>pub\s+)?static\s+ref\s+(?P<name>[a-zA-Z_][a-zA-Z0-9_]*):\s*(?P<type>[^=]+)'
                for match in re.finditer(static_pattern, macro_content):
                    # 提取完整声明直到等号
                    content = match.group(0).strip()
                    if not content.endswith(';'):
                        content += ';'

                    declarations.append(RustDeclaration(
                        kind="static",
                        name=match.group('name'),
                        visibility=match.group('vis').strip() if match.group('vis') else '',
                        generics=None,
                        line_number=i + 1,
                        content=content
                    ))
                i = end_line + 1
                continue

            # 处理普通静态声明
            static_pattern = r'(?P<vis>pub\s+)?static\s+(?:mut\s+)?(?P<name>[a-zA-Z_][a-zA-Z0-9_]*)\s*:'
            match = re.match(static_pattern, line)
            if match:
                full_decl, end_line = self.get_complete_declaration(i)
                decl_end = full_decl.find('=')
                if decl_end == -1:
                    decl_end = len(full_decl)
                content = full_decl[:decl_end].strip()
                if not content.endswith(';'):
                    content += ';'

                declarations.append(RustDeclaration(
                    kind="static",
                    name=match.group('name'),
                    visibility=match.group('vis').strip() if match.group('vis') else '',
                    generics=None,
                    line_number=i + 1,
                    content=content
                ))
                i = end_line + 1
                continue

            i += 1

        return declarations

    def _find_declarations(self, pattern: str, kind: str) -> List[RustDeclaration]:
        """通用的声明查找方法"""
        declarations = []
        i = 0
        while i < len(self.lines):
            line = self.lines[i].strip()
            if not line or line.startswith('//'):
                i += 1
                continue

            match = re.search(pattern, line)
            if match:
                # 获取完整声明
                full_decl, end_line = self.get_complete_declaration(i)

                # 提取声明部分（到大括号之前）
                decl_end = full_decl.find('{')
                if decl_end == -1:
                    decl_end = len(full_decl)
                content = full_decl[:decl_end].strip()
                if not content.endswith(';'):
                    content += ';'

                declarations.append(RustDeclaration(
                    kind=kind,
                    name=match.group('name'),
                    visibility=match.group('vis').strip() if match.group('vis') else '',
                    generics=match.group('generics') if 'generics' in match.groupdict() else None,
                    line_number=i + 1,
                    content=content
                ))

                i = end_line + 1
                continue

            i += 1

        return declarations

# 使用示例
if __name__ == '__main__':
    # 示例 Rust 代码
    sample_code = '''
pub struct Point<T>;

pub struct Data<T> {
    field: T,
}

pub fn calculate_distance<T>(p1: Point<T>, p2: Point<T>) -> T;

pub fn process(data: &str) {
    // 实现部分
}

pub union IntOrFloat {
    i: i32,
    f: f32,
}

pub enum Direction {
    North,
    South,
}

pub static GLOBAL_COUNTER: AtomicUsize = AtomicUsize::new(0);

fn insert(&mut self, 
          value: &[u8]) {
    for i in 0..self.num_functions {
        let hash = murmur3(value, SALTS[i]);
        let index = hash as usize % self.size;
        self.set_bit(index);
    }
}
'''

    test_code = '''
// 常规静态方法
impl BloomFilter {
    pub fn new(size: usize, num_hash_functions: usize) -> Self {
        // 实现
    }
    
    fn calculate_optimal_size(expected_elements: usize, false_positive_rate: f64) -> usize {
        // 实现
    }
}

// 泛型实现中的静态方法
impl<T> Queue<T> {
    pub fn with_capacity(capacity: usize) -> Self {
        // 实现
    }
    
    pub fn default() -> Self where T: Default {
        // 实现
    }
}

// trait 实现中的静态方法
impl From<Vec<u8>> for BloomFilter {
    fn from(data: Vec<u8>) -> Self {
        // 实现
    }
}

// 复杂的静态方法
impl BloomFilter {
    pub fn create_optimal(
        expected_elements: usize,
        false_positive_rate: f64,
        hash_functions: Vec<Box<dyn Fn(&[u8]) -> u64>>
    ) -> Result<Self, Error> {
        // 实现
    }
}

// 异步静态方法
impl AsyncProcessor {
    pub async fn initialize(config: Config) -> Result<Self, InitError> {
        // 实现
    }
    
    pub async fn create_and_process(
        input: Vec<u8>,
        options: ProcessOptions
    ) -> Result<ProcessedData, ProcessError> {
        // 实现
    }
}
'''
    line_code = "pub fn process(data: &str) {"

    error_code = """lazy_static! {
    pub static ref TOTAL_SHAPES: Mutex<i32> = Mutex::new(0);
}
"""

    for code in [sample_code, test_code, line_code, error_code]:
        parser = RustParser(code)
        declarations = parser.parse_all()

        print("\n找到的声明:")
        for decl in declarations:
            print(f"\n{decl.kind} '{decl.name}' at line {decl.line_number}")
            print(f"Visibility: {decl.visibility}")
            if decl.generics:
                print(f"Generics: {decl.generics}")
            print(f"Content: {decl.content}")
